/**
 * @file KDTree.h
 * @brief KD树实现，用于高效的最近邻搜索
 * 
 * 该类实现了一个优化的KD树数据结构，支持：
 * 1. K近邻搜索
 * 2. 八象限搜索
 * 3. SIMD优化的距离计算
 * 4. 内存对齐优化
 */
#pragma once
#include <vector>
#include <utility>
#include "Point3D.h"

class KDTree {
public:
    /**
     * @brief KD树节点结构
     * 
     * 使用32字节对齐以优化内存访问和SIMD计算
     */
    struct alignas(32) KDNode {
        Point3D point;           ///< 节点的空间坐标
        double s1_value;         ///< 节点的S1值
        int axis;               ///< 分割轴(0=x, 1=y, 2=z)
        KDNode *left, *right;   ///< 左右子节点指针

        /**
         * @brief 节点构造函数
         * @param p 空间坐标
         * @param s1 S1值
         */
        KDNode(const Point3D& p, double s1)
            : point(p), s1_value(s1),
              axis(0), left(nullptr), right(nullptr) {}
        
        // 禁用拷贝操作
        KDNode(const KDNode&) = delete;
        KDNode& operator=(const KDNode&) = delete;
        
        ~KDNode() = default;
    };

    /**
     * @brief 默认构造函数
     */
    KDTree();

    /**
     * @brief 析构函数，释放树中所有节点
     */
    ~KDTree();
    
    /**
     * @brief 构建KD树
     * @param points 输入点集，包含坐标和S1值
     * 
     * 使用递归方法构建平衡的KD树
     */
    void build(std::vector<std::pair<Point3D, double>>& points);

    /**
     * @brief 查找K个最近邻点
     * @param target 目标点
     * @param k 需要查找的邻居数量
     * @return 返回K个最近邻点及其距离
     */
    std::vector<std::pair<double, KDNode*>> findKNearest(const Point3D& target, int k) const;

    /**
     * @brief 在指定象限中查找最近点
     * @param target 目标点
     * @param x_dir X方向(1=正,-1=负)
     * @param y_dir Y方向(1=正,-1=负)
     * @param z_dir Z方向(1=正,-1=负)
     * @return 返回最近点及其距离
     */
    std::pair<double, KDNode*> findNearestInOctant(const Point3D& target, 
        int x_dir, int y_dir, int z_dir) const;

    /**
     * @brief 在八个象限中查找最近点
     * @param target 目标点
     * @return 返回八个象限中的最近点及其距离
     */
    std::vector<std::pair<double, KDNode *>> findOctantNeighbors(const Point3D &target) const;

    /**
     * @brief K近邻搜索的另一个重载版本
     * @param target 目标点
     * @param k 需要查找的邻居数量
     * @param nearest 输出参数，存储查找结果
     */
    void findKNearest(const Point3D &target, int k, 
                     std::vector<std::pair<double, KDNode*>> &nearest) const;

private:
    KDNode* root;  ///< KD树的根节点

    /**
     * @brief 递归构建KD树
     * @param points 输入点集
     * @param depth 当前深度
     * @param start 点集起始索引
     * @param end 点集结束索引
     * @return 返回构建的子树根节点
     */
    KDNode *buildTree(std::vector<std::pair<Point3D, double>> &points,
                      int depth, int start, int end);
    
    /**
     * @brief 递归查找最近邻点
     * @param node 当前节点
     * @param target 目标点
     * @param nearest 当前找到的最近点集
     * @param max_dist 当前最大距离
     * @param k 需要查找的邻居数量
     */
    void findNearestNeighbors(KDNode *node, const Point3D &target,
                             std::vector<std::pair<double, KDNode *>> &nearest,
                             double &max_dist, int k) const;
    
    /**
     * @brief 在指定象限中递归查找最近点
     * @param node 当前节点
     * @param target 目标点
     * @param nearest 当前找到的最近点集
     * @param max_dist 当前最大距离
     * @param k 需要查找的邻居数量
     * @param x_dir X方向限制
     * @param y_dir Y方向限制
     * @param z_dir Z方向限制
     */
    void findNearestNeighborsInOctant(KDNode *node, const Point3D &target,
                                     std::vector<std::pair<double, KDNode *>> &nearest,
                                     double &max_dist, int k,
                                     int x_dir, int y_dir, int z_dir) const;
    
    /**
     * @brief 清理KD树节点
     * @param node 要清理的节点
     */
    void cleanup(KDNode *node);

    /**
     * @brief 计算两点间的平方距离
     * @param p1 第一个点
     * @param p2 第二个点
     * @return 返回两点间距离的平方
     */
    static double calculateDistanceSquared(const Point3D &p1, const Point3D &p2);
}; 